"""
MCP客户端核心实现

负责WebSocket连接管理、消息收发、心跳维护等核心功能
"""

import asyncio
import uuid
from typing import Optional, Dict, Any, Callable
from datetime import datetime
import websockets
from websockets.exceptions import ConnectionClosed, WebSocketException
import json

from app.core.logging_config import app_logger
from app.core.config import settings
from .message_handler import MCPMessageHandler
from .exceptions import (
    MCPConnectionError, 
    MCPTimeoutError, 
    MCPAuthenticationError,
    MCPMessageError
)


class MCPClient:
    """
    MCP客户端核心类
    
    功能：
    1. WebSocket连接建立与管理
    2. 消息发送与接收
    3. 心跳维护
    4. 自动重连
    5. 连接状态监控
    
    使用示例：
    ```python
    client = MCPClient(server_url="ws://localhost:8080/mcp")
    await client.connect()
    response = await client.send_request({"query": "test"})
    await client.disconnect()
    ```
    """
    
    def __init__(
        self,
        server_url: str,
        session_id: Optional[str] = None,
        timeout: int = None,
        max_retries: int = None,
        heartbeat_interval: int = None,
        on_message: Optional[Callable] = None,
        on_error: Optional[Callable] = None,
        on_connect: Optional[Callable] = None,
        on_disconnect: Optional[Callable] = None
    ):
        """
        初始化MCP客户端
        
        Args:
            server_url: MCP服务器WebSocket地址
            session_id: 会话ID（可选，默认自动生成）
            timeout: 连接超时时间（秒）
            max_retries: 最大重试次数
            heartbeat_interval: 心跳间隔（秒）
            on_message: 消息接收回调函数
            on_error: 错误处理回调函数
            on_connect: 连接建立回调函数
            on_disconnect: 断开连接回调函数
        """
        self.server_url = server_url
        self.session_id = session_id or str(uuid.uuid4())
        self.timeout = timeout or settings.mcp_connection_timeout
        self.max_retries = max_retries or settings.mcp_max_retries
        self.heartbeat_interval = heartbeat_interval or settings.mcp_heartbeat_interval
        
        # WebSocket连接对象
        self.websocket: Optional[websockets.WebSocketClientProtocol] = None
        
        # 消息处理器
        self.message_handler = MCPMessageHandler()
        
        # 连接状态
        self.is_connected = False
        self.is_connecting = False
        self.retry_count = 0
        
        # 回调函数
        self.on_message = on_message
        self.on_error = on_error
        self.on_connect = on_connect
        self.on_disconnect = on_disconnect
        
        # 异步任务
        self._heartbeat_task: Optional[asyncio.Task] = None
        self._receive_task: Optional[asyncio.Task] = None
        
        # 消息队列（用于请求-响应匹配）
        self._pending_requests: Dict[str, asyncio.Future] = {}
        
        # 日志记录器
        self.logger = app_logger
        
    async def connect(self) -> bool:
        """
        建立WebSocket连接
        
        执行流程：
        1. 检查是否已连接
        2. 建立WebSocket连接
        3. 发送CONNECT消息
        4. 等待CONNECT_ACK响应
        5. 启动心跳和接收任务
        
        Returns:
            bool: 连接是否成功
            
        Raises:
            MCPConnectionError: 连接失败
            MCPTimeoutError: 连接超时
        """
        if self.is_connected:
            self.logger.warning(f"MCP客户端已连接，session_id={self.session_id}")
            return True
            
        if self.is_connecting:
            self.logger.warning(f"MCP客户端正在连接中，session_id={self.session_id}")
            return False
            
        self.is_connecting = True
        self.logger.info(f"开始连接MCP服务器: {self.server_url}, session_id={self.session_id}")
        
        try:
            # 建立WebSocket连接
            self.websocket = await asyncio.wait_for(
                websockets.connect(
                    self.server_url,
                    ping_interval=None,  # 禁用默认ping，使用自定义心跳
                    close_timeout=10
                ),
                timeout=self.timeout
            )
            
            # 等待一小段时间让服务器准备就绪
            await asyncio.sleep(0.1)
            
            # 发送CONNECT消息
            connect_message = self.message_handler.create_connect_message(
                session_id=self.session_id
            )
            await self._send_raw_message(connect_message)
            
            # 等待CONNECT_ACK响应
            ack_message = await asyncio.wait_for(
                self._receive_raw_message(),
                timeout=self.timeout
            )
            
            # 验证响应
            if not self.message_handler.is_connect_ack(ack_message):
                raise MCPConnectionError("未收到有效的CONNECT_ACK响应")
            
            # 连接成功
            self.is_connected = True
            self.is_connecting = False
            self.retry_count = 0
            
            self.logger.info(f"MCP连接建立成功, session_id={self.session_id}")
            
            # 启动后台任务
            self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
            self._receive_task = asyncio.create_task(self._receive_loop())
            
            # 触发连接成功回调
            if self.on_connect:
                if asyncio.iscoroutinefunction(self.on_connect):
                    await self.on_connect(self.session_id)
                else:
                    self.on_connect(self.session_id)
            
            return True
            
        except asyncio.TimeoutError:
            self.is_connecting = False
            error_msg = f"MCP连接超时: {self.server_url}"
            self.logger.error(error_msg)
            raise MCPTimeoutError(error_msg)
            
        except Exception as e:
            self.is_connecting = False
            error_msg = f"MCP连接失败: {str(e)}"
            self.logger.error(error_msg, exc_info=e)
            raise MCPConnectionError(error_msg) from e
    
    async def disconnect(self) -> None:
        """
        断开WebSocket连接
        
        执行流程：
        1. 发送DISCONNECT消息
        2. 取消心跳和接收任务
        3. 关闭WebSocket连接
        4. 清理资源
        """
        if not self.is_connected:
            self.logger.warning(f"MCP客户端未连接，无需断开，session_id={self.session_id}")
            return
        
        self.logger.info(f"开始断开MCP连接, session_id={self.session_id}")
        
        try:
            # 发送DISCONNECT消息
            disconnect_message = self.message_handler.create_disconnect_message(
                session_id=self.session_id
            )
            await self._send_raw_message(disconnect_message)
            
        except Exception as e:
            self.logger.error(f"发送DISCONNECT消息失败: {str(e)}")
        
        finally:
            # 取消后台任务
            if self._heartbeat_task:
                self._heartbeat_task.cancel()
                try:
                    await self._heartbeat_task
                except asyncio.CancelledError:
                    pass
                    
            if self._receive_task:
                self._receive_task.cancel()
                try:
                    await self._receive_task
                except asyncio.CancelledError:
                    pass
            
            # 关闭WebSocket连接
            if self.websocket:
                await self.websocket.close()
            
            # 清理状态
            self.is_connected = False
            self.websocket = None
            
            # 清理待处理请求
            for future in self._pending_requests.values():
                if not future.done():
                    future.set_exception(MCPConnectionError("连接已断开"))
            self._pending_requests.clear()
            
            self.logger.info(f"MCP连接已断开, session_id={self.session_id}")
            
            # 触发断开连接回调
            if self.on_disconnect:
                if asyncio.iscoroutinefunction(self.on_disconnect):
                    await self.on_disconnect(self.session_id)
                else:
                    self.on_disconnect(self.session_id)
    
    async def send_request(
        self, 
        payload: Dict[str, Any],
        timeout: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        发送REQUEST消息并等待RESPONSE响应
        
        Args:
            payload: 请求负载数据
            timeout: 超时时间（秒），默认使用客户端配置
            
        Returns:
            Dict[str, Any]: 响应负载数据
            
        Raises:
            MCPConnectionError: 连接未建立
            MCPTimeoutError: 请求超时
            MCPMessageError: 消息发送失败
        """
        if not self.is_connected:
            raise MCPConnectionError("MCP连接未建立")
        
        # 创建REQUEST消息
        message_id = str(uuid.uuid4())
        request_message = self.message_handler.create_request_message(
            session_id=self.session_id,
            message_id=message_id,
            payload=payload
        )
        
        # 创建Future用于等待响应
        future = asyncio.Future()
        self._pending_requests[message_id] = future
        
        try:
            # 发送消息
            await self._send_raw_message(request_message)
            self.logger.debug(f"发送REQUEST消息, message_id={message_id}")
            
            # 等待响应
            timeout_value = timeout or settings.mcp_request_timeout
            response = await asyncio.wait_for(future, timeout=timeout_value)
            
            self.logger.debug(f"收到RESPONSE响应, message_id={message_id}")
            return response
            
        except asyncio.TimeoutError:
            error_msg = f"请求超时, message_id={message_id}"
            self.logger.error(error_msg)
            raise MCPTimeoutError(error_msg)
            
        except Exception as e:
            error_msg = f"请求失败: {str(e)}"
            self.logger.error(error_msg, exc_info=e)
            raise MCPMessageError(error_msg) from e
            
        finally:
            # 清理待处理请求
            self._pending_requests.pop(message_id, None)
    
    async def _send_raw_message(self, message: Dict[str, Any]) -> None:
        """发送原始消息（内部方法）"""
        if not self.websocket:
            raise MCPConnectionError("WebSocket连接不存在")
        
        try:
            await self.websocket.send(json.dumps(message))
        except Exception as e:
            raise MCPMessageError(f"消息发送失败: {str(e)}") from e
    
    async def _receive_raw_message(self) -> Dict[str, Any]:
        """接收原始消息（内部方法）"""
        if not self.websocket:
            raise MCPConnectionError("WebSocket连接不存在")
        
        try:
            raw_data = await self.websocket.recv()
            return json.loads(raw_data)
        except Exception as e:
            raise MCPMessageError(f"消息接收失败: {str(e)}") from e
    
    async def _heartbeat_loop(self) -> None:
        """心跳循环（后台任务）"""
        self.logger.info(f"心跳任务启动, interval={self.heartbeat_interval}s")
        
        try:
            while self.is_connected:
                await asyncio.sleep(self.heartbeat_interval)
                
                if not self.is_connected:
                    break
                
                try:
                    # 发送PING消息
                    ping_message = self.message_handler.create_ping_message(
                        session_id=self.session_id
                    )
                    await self._send_raw_message(ping_message)
                    self.logger.debug("发送PING心跳")
                    
                except Exception as e:
                    self.logger.error(f"心跳发送失败: {str(e)}")
                    # 心跳失败可能意味着连接断开，触发重连
                    await self._handle_connection_lost()
                    break
                    
        except asyncio.CancelledError:
            self.logger.info("心跳任务已取消")
        except Exception as e:
            self.logger.error(f"心跳任务异常: {str(e)}", exc_info=e)
    
    async def _receive_loop(self) -> None:
        """消息接收循环（后台任务）"""
        self.logger.info("消息接收任务启动")
        
        try:
            while self.is_connected:
                try:
                    # 接收消息
                    message = await self._receive_raw_message()
                    await self._handle_message(message)
                    
                except ConnectionClosed:
                    self.logger.warning("WebSocket连接已关闭")
                    await self._handle_connection_lost()
                    break
                    
                except Exception as e:
                    self.logger.error(f"消息接收异常: {str(e)}", exc_info=e)
                    
        except asyncio.CancelledError:
            self.logger.info("消息接收任务已取消")
        except Exception as e:
            self.logger.error(f"消息接收任务异常: {str(e)}", exc_info=e)
    
    async def _handle_message(self, message: Dict[str, Any]) -> None:
        """处理接收到的消息（内部方法）"""
        try:
            message_type = message.get("type")
            message_id = message.get("message_id")
            
            self.logger.debug(f"收到消息: type={message_type}, message_id={message_id}")
            
            # 处理RESPONSE消息
            if message_type == "RESPONSE":
                if message_id in self._pending_requests:
                    future = self._pending_requests[message_id]
                    if not future.done():
                        future.set_result(message.get("payload", {}))
            
            # 处理PONG消息
            elif message_type == "PONG":
                self.logger.debug("收到PONG心跳响应")
            
            # 处理ERROR消息
            elif message_type == "ERROR":
                error_code, error_message = self.message_handler.extract_error(message)
                self.logger.error(f"收到ERROR消息: code={error_code}, message={error_message}")
                
                # 如果有对应的请求，设置异常
                if message_id in self._pending_requests:
                    future = self._pending_requests[message_id]
                    if not future.done():
                        future.set_exception(MCPMessageError(f"{error_code}: {error_message}"))
                
                # 触发错误回调
                if self.on_error:
                    if asyncio.iscoroutinefunction(self.on_error):
                        await self.on_error(error_code, error_message)
                    else:
                        self.on_error(error_code, error_message)
            
            # 处理PING消息（服务器主动发送）
            elif message_type == "PING":
                # 回复PONG
                pong_message = self.message_handler.create_pong_message(
                    session_id=self.session_id
                )
                await self._send_raw_message(pong_message)
                self.logger.debug("回复PONG心跳")
            
            # 触发消息接收回调
            if self.on_message:
                if asyncio.iscoroutinefunction(self.on_message):
                    await self.on_message(message)
                else:
                    self.on_message(message)
                
        except Exception as e:
            self.logger.error(f"消息处理异常: {str(e)}", exc_info=e)
    
    async def _handle_connection_lost(self) -> None:
        """处理连接丢失（内部方法）"""
        self.logger.warning("检测到连接丢失")
        
        # 更新状态
        self.is_connected = False
        
        # 尝试重连
        if self.retry_count < self.max_retries:
            self.retry_count += 1
            self.logger.info(f"尝试重连 ({self.retry_count}/{self.max_retries})")
            
            try:
                await asyncio.sleep(5)  # 等待5秒后重连
                await self.connect()
            except Exception as e:
                self.logger.error(f"重连失败: {str(e)}")
        else:
            self.logger.error(f"已达到最大重试次数 ({self.max_retries})，放弃重连")

